import wave
import torch
import numpy as np

def read_wav(file_path, return_array=False):
    with wave.open(file_path, "rb") as f:
        n_frames = f.getnframes()
        stream = f.readframes(n_frames)
    if not return_array:
        return stream
    else:
        array = np.frombuffer(stream, dtype=np.int16)
        return array

def write_wave(file_path, pcm):
    f = wave.open(file_path, "wb")
    # 配置声道数，量化位和取样数
    f.setnchannels(1)
    f.setsampwidth(2)
    f.setframerate(16000)

    # 将wave转为二进制写入文件
    f.writeframes(pcm)
    f.close()

def int2float(sound):
    abs_max = np.abs(sound).max()
    sound = sound.astype("float32")
    if abs_max >0:
        sound *= 1/32768
    sound = sound.squeeze()
    return sound

def float2int(sound):
    sound *= 32768
    sound = sound.astype("int16")
    return sound

def bytes2tensor(pcm:bytes) -> torch.Tensor:
    waveform_i16 = np.frombuffer(pcm, dtype=np.int16)
    waveform_i32 = int2float(waveform_i16)
    tensor = torch.tensor(waveform_i32, dtype=torch.float32)
    tensor = tensor.unsqueeze(0)
    return tensor

def copy_stream(stream, repeat_times):
    tmp_streams = [stream for _ in range(repeat_times)]
    return b"".join(tmp_streams)
